import tensorflow as tf
from PIL import Image
import numpy as np
import os as os

from src.digits.digit_train import CNN
from src.utils.visualize import *

'''
python 3.7
tensorflow 2.0.0b0
pillow(PIL) 4.3.0
'''


class Predict(object):

    def __init__(self):
        latest = tf.train.latest_checkpoint('./ckpt')
        self.cnn = CNN()
        # 恢复网络权重
        self.cnn.model.load_weights(latest)

    def predict(self, image_path):
        # 以黑白方式读取图片
        img = Image.open(image_path).convert('L')
        # 将图片统一缩放到28*28px
        img = img.resize((28, 28), Image.ANTIALIAS)
        img = np.reshape(img, (28, 28, 1)) / 255.
        x = np.array([1 - img])

        # API refer: https://keras.io/models/model/
        y = self.cnn.model.predict(x)

        # 因为x只传入了一张图片，取y[0]即可
        # np.argmax()取得最大值的下标，即代表的数字
        print(image_path)
        # print(np.argmax(y[0]))
        print('\t-> 这个图片写的是：', get_mapping(np.argmax(y[0])))


if __name__ == "__main__":
    image_dir = '../../number_images/'
    # cwd = os.getcwd()
    files = os.listdir(image_dir)
    # 遍历批量导入预测文件夹内的所有图片
    app = Predict()
    for image in files:
        app.predict(image_dir + image)
